'''
Copyright (c) 2013 Andrew Meyer

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
associated documentation files (the "Software"), to deal in the Software without restriction,
including without limitation the rights to use, copy, modify, merge, publish, distribute,
sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial
portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES
OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
'''


bl_info = {
	"name": "Ascend Exporter",
	"author": "Andrew Meyer",
	"version": (2, 3, 0),
	"blender": (2, 6, 8),
	"location": "File > Import-Export",
	"description": "Export to Ascend model format",
	"category": "Import-Export",
	"wiki_url": "ascend3d.codeplex.com/documentation",
	"tracker_url": "ascend3d.codeplex.com/workitem/list/basic"}

import bpy, os
from xml.etree.ElementTree import Element, ElementTree, tostring
from xml.dom.minidom import parseString
from bpy.types import Operator, PropertyGroup
from bpy.props import StringProperty, BoolProperty, EnumProperty, IntProperty, CollectionProperty
from bpy_extras.io_utils import ExportHelper
from mathutils import Vector

#
#	Utility functions
#
def get_armature_skin(arm, scene):
	for obj in scene.objects:
		if obj.find_armature() == arm:
			return obj
	
	return None


def is_animated(obj):
	return obj is not None and obj.animation_data is not None and \
		obj.animation_data.action is not None


#
#	Main exporter
#
class AscendExporter(Operator, ExportHelper):
	bl_idname = "ascend.exporter"
	bl_label = "Export Ascend"
	filename_ext = ".asc"
	
	filter_glob = StringProperty(default="*.asc", options={'HIDDEN'})
	
	def execute(self, context):
		try:
			if context.mode != 'OBJECT':
				bpy.ops.object.mode_set(mode='OBJECT')
			
			# get tree root if required
			if context.scene.export_hierarchy:
				obj = self.get_ancestor(context.object, context.scene)
			else:
				obj = context.object
			
			if obj.type != 'MESH':
				raise Exception("Cannot export top-level object of type %s" % obj.type)
			
			# begin exporting
			curr = context.scene.frame_current
			context.scene.frame_set(context.scene.frame_start)
			ver = bl_info["version"]
			self.tree = ElementTree(Element("Ascend",
				{"ExporterVersion" : str(ver[0]) + "." + str(ver[1]) + "." + str(ver[2])}))
			self.write_mesh(obj, context.scene, self.tree.getroot())
			context.scene.frame_set(curr)
			
			# write to file
			self.write_xml()
		
		except Exception as ex:
			bpy.ops.ascend.message_box('INVOKE_DEFAULT', err=True, message=str(ex))
			return {'FINISHED'}
		
		bpy.ops.ascend.message_box('INVOKE_DEFAULT', message="Export Successful!")
		return {'FINISHED'}
	
	@classmethod
	def poll(self, context):
		obj = context.object
		
		# object must be a mesh or an armature with a valid skin
		return obj is not None and \
			(obj.type == 'MESH' or \
			(obj.type == 'ARMATURE' and get_armature_skin(obj, context.scene) is not None))
	
	def get_ancestor(self, obj, scene):
		# if object is an armature, try to find its skin
		if obj.type == 'ARMATURE':
			skin = get_armature_skin(obj, scene)
			
			if skin is not None:
				return skin
		
		if obj.parent is None:
			return obj
		else:
			return self.get_ancestor(obj.parent, scene)
	
	def write_mesh(self, mesh, scene, xmlparent, parent_bone = None):
		# ensure tesselation faces
		mesh.data.calc_tessface()
		
		# decide node type
		arm = mesh.find_armature()
		shouldWriteAnim = scene.export_anim and (is_animated(mesh) or is_animated(arm))
		shouldWriteArm = scene.export_arm and mesh.find_armature() is not None
		node = \
			Element("AnimatedSceneNode" if shouldWriteAnim or shouldWriteArm else "ModelSceneNode")
		xmlparent.append(node)
		
		# write node properties
		node.set("Name", mesh.name)
		
		shouldWriteUVs = scene.export_uv and mesh.data.tessface_uv_textures.active is not None and \
			len(mesh.data.tessface_uv_textures.active.data) > 0
		if shouldWriteUVs and mesh.data.tessface_uv_textures.active.data[0].image is not None:
			node.set("Texture",
				bpy.path.basename(mesh.data.tessface_uv_textures.active.data[0].image.filepath))
		
		if len(mesh.data.materials) > 0:
			mtl = mesh.data.materials[0]
			alpha = mtl.alpha if mtl.use_transparency else 1.0
			node.set("Color", "%f %f %f %f" %
				(mtl.diffuse_color.r, mtl.diffuse_color.g, mtl.diffuse_color.b, alpha))
		else:
			node.set("Color", "1.0 1.0 1.0 1.0")
		
		# write transform
		if parent_bone is not None:
			mat = parent_bone.matrix.inverted() * mesh.matrix_local
		else:
			mat = self.get_mesh_transform(mesh, scene)
		
		self.write_transform(mat, node)
		
		# write vertices
		self.write_vertices(mesh.data.vertices, scene, node)
		
		# write faces
		self.write_faces(mesh, scene, shouldWriteUVs, node)
		
		# write animation frames
		if shouldWriteAnim:
			if is_animated(mesh):
				self.write_animation_frames(mesh, scene, node)
			
			# write animation definitions
			anims = Element("Animations")
			for ad in scene.anims_coll:
				if ad.object_enum == mesh.name or ad.object_enum == '~ALL':
					anim = Element("Animation",
						{
							"Name" : ad.name,
							"Start" : str(ad.start - scene.frame_start),
							"End" : str(ad.end - scene.frame_start)
						})
					
					anims.append(anim)
			
			node.append(anims)
		
		# write armature
		if scene.export_arm and mesh.find_armature() is not None:
			self.write_armature(mesh, scene, node)
		
		# write children
		if scene.export_hierarchy:
			self.write_node_children(mesh, scene, node)
		else:
			ch = Element("Children")
			node.append(ch)
	
	def get_mesh_transform(self, mesh, scene):
		# if not exporting hierarchy or parent is an armature, mesh should be at the top of the
		# scene graph
		if not scene.export_hierarchy or \
			(mesh.parent is not None and mesh.parent.type == 'ARMATURE'):
			return mesh.matrix_world
		else:
			return mesh.matrix_local
	
	def write_transform(self, mat, xmlparent):
		pos = mat.to_translation()
		xmlparent.set("Position", "%f %f %f" % (pos.x, pos.z, -pos.y))
		rot = mat.to_euler('ZYX')
		xmlparent.set("Rotation", "%f %f %f" % (rot.x, rot.z, -rot.y))
		scl = mat.to_scale()
		xmlparent.set("Scale", "%f %f %f" % (scl.x, scl.z, scl.y))
	
	def write_vertices(self, vertices, scene, xmlparent):
		verts = Element("Vertices")
		
		for v in vertices:
			vert = Element("Vertex")
			vert.set("Point", "%f %f %f" % (v.co.x, v.co.z, -v.co.y))
			
			# add vertex normals if they might be needed
			if scene.normals_enum != 'NONE' and scene.normals_enum != 'FACE':
				n = v.normal
				vert.set("Normal", "%f %f %f" % (n.x, n.z, -n.y))
			
			verts.append(vert)
		
		xmlparent.append(verts)
		
	def write_faces(self, mesh, scene, shouldWriteUVs, xmlparent):
		faces = Element("Faces")
		
		for f in mesh.data.tessfaces:
			face = Element("Face")
			
			# add face normal if needed for this face
			if scene.normals_enum == 'FACE' or \
				(scene.normals_enum == 'AUTO' and not f.use_smooth):
				n = f.normal
				face.set("Normal", "%f %f %f" % (n.x, n.z, -n.y))
			
			for v in range(len(f.vertices)):
				vertex = Element("FaceVertex")
				idx = f.vertices[v]
				vertex.set("Index", str(idx))
				
				if shouldWriteUVs:
					uv = mesh.data.tessface_uv_textures.active.data[f.index].uv[v]
					vertex.set("UV", "%f %f" % (uv[0], 1 - uv[1]))
				
				face.append(vertex)
			
			faces.append(face)
		
		xmlparent.append(faces)
	
	def write_animation_frames(self, obj, scene, xmlparent):
		frames = Element("Frames")
		
		for f in range(scene.frame_start, scene.frame_end + 1):
			scene.frame_set(f)
			frame = Element("Frame")
			self.write_transform(obj.matrix_local, frame)
			frames.append(frame)
		
		xmlparent.append(frames)
	
	def write_node_children(self, parent, scene, xmlparent):
		ch = Element("Children")
		
		for child in parent.children:
			if child.type == 'MESH':
				self.write_mesh(child, scene, ch)
		
		xmlparent.append(ch)
	
	def write_armature(self, mesh, scene, xmlparent):
		armature = mesh.find_armature()
		
		arm = Element("Armature")
		arm.set("Name", armature.name)
		
		shouldWriteAnim = scene.export_anim and is_animated(armature)
		
		# write bones (beginning with top-level bones)
		rootBones = Element("RootBones")
		for bone in armature.pose.bones:
			if bone.parent is None:
				self.write_bone(bone, scene, mesh, shouldWriteAnim, rootBones)
		
		arm.append(rootBones)
		
		# write attached children
		if scene.export_hierarchy:
			for child in armature.children:
				if child.type == 'MESH' and child.parent_bone is not None and \
					child.parent_bone in armature.data.bones:
					for bn in arm.iter("Bone"):
						if bn.get("Name") == child.parent_bone:
							ch = bn.find("Children")
							self.write_mesh(child, scene, ch,
								armature.pose.bones[child.parent_bone])
							break
		
		xmlparent.append(arm)
	
	def write_bone(self, bone, scene, mesh, shouldWriteAnim, xmlparent):
		bn = Element("Bone")
		bn.set("Name", bone.name)
		
		# write resting transform
		if bone.parent is None:
			mat = bone.bone.matrix_local
		else:
			mat = bone.parent.bone.matrix_local.inverted() * bone.bone.matrix_local
		
		self.write_transform(mat, bn)
		
		# write animation frames
		if shouldWriteAnim:
			self.write_bone_animation_frames(bone, scene, bn)
		
		# write influences
		self.write_bone_influences(bone, mesh, bn)
		
		# write child bones
		ch = Element("Children")
		for child in bone.children:
			self.write_bone(child, scene, mesh, shouldWriteAnim, ch)
		
		bn.append(ch)
		
		xmlparent.append(bn)
	
	def write_bone_animation_frames(self, bone, scene, xmlparent):
		frames = Element("Frames")
		
		for f in range(scene.frame_start, scene.frame_end + 1):
			scene.frame_set(f)
			frame = Element("Frame")
			
			if bone.parent is None:
				mat = bone.matrix
			else:
				mat = bone.parent.matrix.inverted() * bone.matrix
			
			self.write_transform(mat, frame)			
			frames.append(frame)
		
		xmlparent.append(frames)
	
	def write_bone_influences(self, bone, mesh, xmlparent):
		infs = Element("Influences")
		
		for vert in mesh.data.vertices:
			for gp in vert.groups:
				if gp.weight != 0.0 and mesh.vertex_groups[gp.group].name == bone.name:
					inf = Element("Influence")
					inf.set("Vertex", str(vert.index))
					inf.set("Weight", str(gp.weight))
					infs.append(inf)
		
		xmlparent.append(infs)
	
	def write_xml(self):
		xmlStr = tostring(self.tree.getroot())
		reparsed = parseString(xmlStr)
		f = open(self.filepath, 'w')
		reparsed.writexml(f, addindent="\t", newl="\n")
	
	def draw(self, context):
		layout = self.layout
		scene = context.scene
		
		# normals UI
		row = layout.row()
		col = row.column()
		col.label("Normals",
			icon=scene.normal_icons[scene.normal_strings.index(scene.normals_enum)])
		col.prop(scene, "normals_enum", expand=True)
		
		# UVs, armatures, hierarchy UI
		col = row.column()
		col.label("")
		col.prop(scene, "export_uv", icon=('TEXTURE' if scene.export_uv else 'X'))
		col.prop(scene, "export_arm", icon=('ARMATURE_DATA' if scene.export_arm else 'X'))
		col.prop(scene, "export_hierarchy", icon=('OOPS' if scene.export_hierarchy else 'X'))
		
		# animation UI
		row = layout.row(align=True)
		row.prop(scene, "export_anim", icon=('RENDER_ANIMATION' if scene.export_anim else 'X'))
				
		if scene.export_anim:
			row.operator("ascend.add_anim_def", text="", icon='ZOOMIN')
			
			for i in reversed(range(len(scene.anims_coll))):
				anim = scene.anims_coll[i]
				box = layout.box()
				row = box.row(align=True)
				row.operator("ascend.remove_anim_def", text="", icon='PANEL_CLOSE').idx = i
				row.prop(anim, "name", text="")
				row = box.row(align=True)
				row.prop(anim, "object_enum", text="", icon='OBJECT_DATA')
				row.prop(anim, "start")
				row.prop(anim, "end")


#
#	Animation definition property group
#
class AnimDef(PropertyGroup):	 
	name = StringProperty(
		name="Name",
		description="The name of the animation",
		default="untitled")
	
	def get_objects(self, context):
		objs = []
		
		if context.scene.export_hierarchy:
			# return all meshes with animations or animated armatures
			for obj in context.scene.objects:
				if obj.type == 'MESH' and (is_animated(obj) or \
					(context.scene.export_arm and is_animated(obj.find_armature()))):
				   objs.append((obj.name, obj.name, obj.name))
			
			if len(objs) > 1:
				objs.insert(0, ('~ALL', "(All)", "Define animation on all animated objects"))
			elif len(objs) == 0:
				objs = [('~NONE', "(None)", "None")]
		elif context.object.type == 'MESH' and (is_animated(context.object) or \
			 context.scene.export_arm and is_animated(context.object.find_armature())):
			# return selected mesh if it is animated or has an animated armature
			obj = context.object
			objs = [(obj.name, obj.name, obj.name)]
		elif context.object.type == 'ARMATURE' and \
			(is_animated(get_armature_skin(context.object)) or\
			 (context.scene.export_arm and is_animated(context.object))):
			# return selected armature's mesh if it or armature is animated
			obj = get_armature_skin(context.object, context.scene)
			objs = [(obj.name, obj.name, obj.name)]
		else:
			objs = [('~NONE', "(None)", "None")]
		
		return objs
	
	object_enum = EnumProperty(
		name = "Object",
		description="The object on which animations are defined",
		items=get_objects)
	
	def update_start(self, context):
		if self.start < context.scene.frame_start:
			self.start = context.scene.frame_start
		
		if self.start > context.scene.frame_end:
			self.start = context.scene.frame_end
		
		self.update_end(context)
	
	def update_end(self, context):
		if self.end < self.start:
			self.end = self.start
		
		if self.end > context.scene.frame_end:
			self.end = context.scene.frame_end
	
	start = IntProperty(
		subtype='UNSIGNED',
		name="Start",
		description="The first frame of the animation",
		update=update_start)
	
	end = IntProperty(
		subtype='UNSIGNED',
		name="End",
		description="The last frame of the animation",
		update=update_end)

bpy.utils.register_class(AnimDef)


#
#	Scene property definitions
#
bpy.types.Scene.normal_strings = ['AUTO', 'VERTEX', 'FACE', 'NONE']
bpy.types.Scene.normal_icons = ['OBJECT_DATA', 'VERTEXSEL', 'FACESEL', 'X']

bpy.types.Scene.normal_options = [
	(bpy.types.Scene.normal_strings[0], "Auto", "Export normals based on shading mode", 0),
	(bpy.types.Scene.normal_strings[1], "Vertex", "Export all normals using vertex normals", 1),
	(bpy.types.Scene.normal_strings[2], "Face", "Export all normals using face normals",	 2),
	(bpy.types.Scene.normal_strings[3], "None", "Do not export normals", 3)]


bpy.types.Scene.normals_enum = EnumProperty(
	name="Normals",
	description="Export normals",
	items=bpy.types.Scene.normal_options,
	default='AUTO')

bpy.types.Scene.export_uv = BoolProperty(
	name="Textures",
	description="Export texture filename and texture coordinates",
	default=True)

bpy.types.Scene.export_arm = BoolProperty(
	name="Armatures",
	description="Export armatures",
	default=True)

bpy.types.Scene.export_hierarchy = BoolProperty(
	name="Hierarchy",
	description="Export all parents and children of selected object",
	default=True)
	
bpy.types.Scene.export_anim = BoolProperty(
	name="Animations",
	description="Export animations",
	default=True)

bpy.types.Scene.anims_coll = CollectionProperty(
	type=AnimDef,
	name="Animation Definitions",
	description="Animations defined on objects in the scene")


#
#	Add animation button
#
class AddAnimDef(Operator):
	'''Add a new animation definition'''
	bl_idname = "ascend.add_anim_def"
	bl_label = "Add Animation Definition"
	
	def execute(self, context):
		anim = context.scene.anims_coll.add()
		anim.start = context.scene.frame_start
		anim.end = context.scene.frame_end
		return {'FINISHED'}


#
#	Remove animation button
#
class RemoveAnimDef(Operator):
	'''Remove animation definition'''
	bl_idname = "ascend.remove_anim_def"
	bl_label = "Remove Animation Definition"
	
	idx = IntProperty()
	
	def execute(self, context):
		context.scene.anims_coll.remove(self.idx)
		return {'FINISHED'}


#
#	Message box
#
class MessageBox(Operator):
	bl_idname = "ascend.message_box"
	bl_label = "Message Box"
	
	message = StringProperty()
	err = BoolProperty(default=False)
 
	def execute(self, context):
		self.report({'INFO'}, self.message)
		return {'FINISHED'}
 
	def invoke(self, context, event):
		wm = context.window_manager
		return wm.invoke_popup(self, width=(300 if self.err else 125))
 
	def draw(self, context):
		self.layout.label("Ascend Exporter")
		self.layout.label(self.message, icon=('ERROR' if self.err else 'FILE_TICK'))


#
#	Registrations
#
def menu_export(self, context):
	self.layout.operator(AscendExporter.bl_idname, text="Ascend (.asc)")


def register():
	bpy.utils.register_module(__name__)
	bpy.types.INFO_MT_file_export.append(menu_export)
	

def unregister():
	bpy.utils.unregister_module(__name__)
	bpy.types.INFO_MT_file_export.remove(menu_export)
	

if __name__ == "__main__":
	register()
